"""
align face for face feature extraction with 5-point landmarks
"""
import cv2
import numpy as np
from skimage import transform as trans


def get_align_src(image_size=112, mode='arcface'):
    arcface_src = np.expand_dims(
        np.array(
            [
                [38.2946, 51.6963],
                [73.5318, 51.5014],
                [56.0252, 71.7366],
                [41.5493, 92.3655],
                [70.7299, 92.2041],
            ],
            dtype=np.float32,
        ),
        axis=0,
    )

    src1 = np.array(
        [
            [51.642, 50.115],
            [57.617, 49.990],
            [35.740, 69.007],
            [51.157, 89.050],
            [57.025, 89.702],
        ],
        dtype=np.float32,
    )

    # <--left
    src2 = np.array(
        [
            [45.031, 50.118],
            [65.568, 50.872],
            [39.677, 68.111],
            [45.177, 86.190],
            [64.246, 86.758],
        ],
        dtype=np.float32,
    )

    # ---frontal
    src3 = np.array(
        [
            [39.730, 51.138],
            [72.270, 51.138],
            [56.000, 68.493],
            [42.463, 87.010],
            [69.537, 87.010],
        ],
        dtype=np.float32,
    )

    # -->right
    src4 = np.array(
        [
            [46.845, 50.872],
            [67.382, 50.118],
            [72.737, 68.111],
            [48.167, 86.758],
            [67.236, 86.190],
        ],
        dtype=np.float32,
    )

    # -->right profile
    src5 = np.array(
        [
            [54.796, 49.990],
            [60.771, 50.115],
            [76.673, 69.007],
            [55.388, 89.702],
            [61.257, 89.050],
        ],
        dtype=np.float32,
    )

    src = np.array([src1, src2, src3, src4, src5])
    src_map = {112: src, 224: src * 2}

    if mode == 'arcface':
        assert image_size == 112
        return arcface_src
    else:
        assert image_size in [112, 224]
        return src_map[image_size]


class FaceAligner(object):
    def __init__(self, image_size=112, mode='arcface'):
        self.src = get_align_src(image_size, mode)
        self.image_size = image_size

    def estimate_norm(self, lmk):
        assert lmk.shape == (5, 2)
        tform = trans.SimilarityTransform()
        lmk_tran = np.insert(lmk, 2, values=np.ones(5), axis=1)
        min_M = []
        min_index = []
        min_error = float("inf")

        # find the source with minimum reprojection error
        for i in np.arange(self.src.shape[0]):
            tform.estimate(lmk, self.src[i])
            M = tform.params[0:2, :]
            results = np.dot(M, lmk_tran.T)
            results = results.T
            error = np.sum(np.sqrt(np.sum((results - self.src[i]) ** 2, axis=1)))
            #         print(error)
            if error < min_error:
                min_error = error
                min_M = M
                min_index = i
            
        return min_M, min_index

    def norm_crop(self, image, landmark):
        M, pose_index = self.estimate_norm(landmark)
        warped = cv2.warpAffine(image, M, (self.image_size, self.image_size), borderValue=0.0)
        return warped


if __name__ == '__main__':
    pass
